package cn.maiyar.common.util;

import lombok.Data;
import lombok.ToString;
import org.xbill.DNS.*;

import java.net.UnknownHostException;

/**
 * dns解析工具类
 *
 * @author c.c
 */
public class DNSUtil {

    public static final String TYPE_A = "A";

    public static final String TYPE_CNAME = "CNAME";

    public static Resolver DEFAULT_RESOLVER = Lookup.getDefaultResolver();

//    public static Resolver HTTP_RESOLVER = new DohResolver("https://dns.google/dns-query");

    public static Resolver HTTP_RESOLVER;

    static {
        try {
            HTTP_RESOLVER = new SimpleResolver("114.114.114.114");
        } catch (UnknownHostException e) {
            e.printStackTrace();
        }
    }

    /**
     * 将输入文本按照A类型进行解析
     *
     * @param name
     * @return
     * @throws TextParseException
     */
    public static DNSRecord resolveOfTypeA(String name) throws TextParseException {
        Lookup lookup = new Lookup(name, Type.A);
        lookup.setResolver(HTTP_RESOLVER);
        Record[] records = null;
        try {
            records = lookup.run();
        } catch (Exception e) {
            lookup.setResolver(DEFAULT_RESOLVER);
            records = lookup.run();
        }
        if (records == null || records.length == 0) {
            return DNSRecord.fail(name, TYPE_A);
        }
        ARecord mx = (ARecord) records[0];
        if (mx.getAddress() == null) {
            return DNSRecord.fail(name, TYPE_A);
        }
        DNSRecord success = DNSRecord.success(name, TYPE_A);
        success.setTtl(mx.getTTL());
        success.setResult(mx.getAddress().getHostAddress());
        return success;
    }

    /**
     * 将输入文本按照CNAME类型进行解析
     *
     * @param name
     * @return
     * @throws TextParseException
     */
    public static DNSRecord resolveOfTypeCName(String name) throws TextParseException {
        Lookup lookup = new Lookup(name, Type.CNAME);
        lookup.setResolver(HTTP_RESOLVER);
        Record[] records = null;
        try {
            records = lookup.run();
        } catch (Exception e) {
            lookup.setResolver(DEFAULT_RESOLVER);
            records = lookup.run();
        }
        if (records == null || records.length == 0) {
            return DNSRecord.fail(name, TYPE_CNAME);
        }
        CNAMERecord mx = (CNAMERecord) records[0];
        if (mx.getTarget() == null) {
            return DNSRecord.fail(name, TYPE_CNAME);
        }
        DNSRecord success = DNSRecord.success(name, TYPE_CNAME);
        success.setTtl(mx.getTTL());
        success.setResult(mx.getTarget().toString(true));
        return success;
    }

    @Data
    @ToString
    public static class DNSRecord {
        // 解析是否成功
        private Boolean success;
        // 输入文本
        private String name;
        // 解析结果
        private String result;
        // ttl时间
        private long ttl;
        // 解析类型
        private String type;

        public static DNSRecord fail(String name, String type) {
            DNSRecord result = new DNSRecord();
            result.setSuccess(false);
            result.setName(name);
            result.setType(type);
            return result;
        }

        public static DNSRecord success(String name, String type) {
            DNSRecord result = new DNSRecord();
            result.setSuccess(true);
            result.setName(name);
            result.setType(type);
            return result;
        }
    }
}
